import cv2
import numpy as np
from PIL import Image

def cvtColor(image):
    '''
    将图像转换成RGB图像，防止灰度图在预测时报错。
    代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
    '''
    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
        return image 
    else:
        image = image.convert('RGB')
        return image 

#---------------------------------------------------#
#   对输入图像进行resize
#---------------------------------------------------#
def resize_image(image, size, letterbox_image):
    if isinstance(image, Image.Image):
        iw, ih  = image.size
        w, h    = size
        if letterbox_image:
            scale   = min(w/iw, h/ih)
            nw      = int(iw*scale)
            nh      = int(ih*scale)

            image   = image.resize((nw,nh), Image.BICUBIC)  #被粘贴的图
            paste_img = Image.new('RGB', size, (128,128,128))   #主图
            paste_img.paste(image, ((w-nw)//2, (h-nh)//2))
        else:
            paste_img = image.resize((w, h), Image.BICUBIC)
    elif type(image) == np.ndarray:
        ih,iw = image.shape[:2]
        h,w = size[:]
        if letterbox_image:
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)

            # 先求左上角
            paste_left_x = (w - nw) // 2
            paste_left_y = (h - nh) // 2
            image = cv2.resize(image,(nh,nw),interpolation=cv2.INTER_CUBIC)
            paste_img = np.ones((h,w,3))*128    #灰度图
            #粘贴
            paste_img[paste_left_y:paste_left_y+nh,paste_left_x:paste_left_x+nw,:] = image[...]
        else:
            paste_img = cv2.resize(image,(h,w),interpolation=cv2.INTER_CUBIC)
    return paste_img

#---------------------------------------------------#
#   获得类
#---------------------------------------------------#
def get_classes(classes_path):
    with open(classes_path, encoding='utf-8') as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names, len(class_names)

#---------------------------------------------------#
#   获得学习率
#---------------------------------------------------#
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
